import os
import sys
import torch
from modelscope import snapshot_download
from diffsynth import ModelManager, FluxImagePipeline

try:
    # 设置模型路径
    model_base_path = 'E:/workspace/llm/text2image/models'
    model_name = 'yiwanji/FLUX_xiao_hong_shu_ji_zhi_zhen_shi_V2'
    model_path = os.path.join(model_base_path, model_name)

    # 如果模型不存在，则下载
    if not os.path.exists(model_path):
        print(f"Downloading model to {model_path}")
        model_dir = snapshot_download(model_name, cache_dir=model_base_path)
        print(f"Model downloaded to {model_dir}")
    else:
        print(f"Using existing model at {model_path}")

    # 设置推理计算精度
    model_manager = ModelManager(torch_dtype=torch.float16)

    # 加载模型 - 使用正确的参数名称 file_path_list 而不是 model_type
    print("Loading model...")
    model_manager.load_models(
        file_path_list=[model_path],  # 正确的参数名称
        torch_dtype=torch.float16,
        device="cuda" if torch.cuda.is_available() else "cpu"
    )

    # 创建pipeline
    print("Creating pipeline...")
    pipe = FluxImagePipeline.from_model_manager(
        model_manager,
        device="cuda" if torch.cuda.is_available() else "cpu",
        torch_dtype=torch.float16
    )

    if pipe is None:
        raise RuntimeError("Failed to create pipeline")

    # 开启CPU卸载以节省显存
    print("Enabling CPU offload...")
    pipe.enable_cpu_offload()

    # 启用量化以进一步减少显存使用
    print("Enabling quantization...")
    if hasattr(pipe, 'dit') and pipe.dit is not None:
        pipe.dit.quantize()
    else:
        print("Warning: DIT quantization not available")

    # 生成图像
    print("Generating image...")
    image = pipe(
        prompt="a beautiful girl",
        negative_prompt="nsfw, nude, bad quality, worst quality",
        seed=0,
        num_inference_steps=30,
        # guidance_scale=7.5
    )

    # 保存图像
    output_dir = "output"
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, "image.jpg")
    image.save(output_path)
    print(f"Image saved to: {output_path}")

except Exception as e:
    print(f"Error occurred: {str(e)}")
    # 打印更详细的错误信息
    import traceback
    traceback.print_exc()
    sys.exit(1)

print("Process completed successfully!")